fix(codegen): emit tensor phis via tensors dict to avoid NameError (#2180) - #2183

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

_{Comment @coderabbitai help to get the list of available commands.}

Review

The diagnosis and the fix direction are right: in this codegen a tensor only ever exists as a tensors["..."] entry (every use site goes through it — e.g. EmitCallToWorker at distributed_codegen.cpp:1018), so a bare x = y Python assignment for a tensor alias is dead on arrival. The new branch in VisitStmt_(AssignStmtPtr) is the correct fix for #2180.

My main comment is that the ~140 new lines in VisitStmt_(IfStmtPtr) can collapse to ~25, and most of the remaining concerns disappear with it.

1. find_yield_source and the phi pre-declaration are redundant

The comment justifying find_yield_source says the then-branch's boundary = zero creates a bare Python name that isn't in the tensors dict. That was true before this PR — but the AssignStmt fix in this same PR removes that case. Walking every way a yield var can be defined:

So by the time VisitStmt(op->then_body_) returns, tensors[<yield var name>] is guaranteed to exist and the yield can be emitted directly, with no RHS back-tracing.

The pre-declaration is redundant for the same reason: ConvertToSSA always synthesizes an else body when it creates phis (convert_to_ssa_pass.cpp:936-939, else_with_yield = make_shared<YieldStmt>(else_yields, ...)), so both branches always assign the phi. That also removes the two fallback initializers, which are the part I'd most like to see gone:

• torch.zeros((1,), dtype=torch.float32) hardcodes shape and dtype, while the phi's own TensorType is right there (compare lines 1051–1070, which derive shape/dtype from the return type).
• tensor_phi_init picks the function's first tensor parameter, which is unrelated to the phi in shape, dtype and meaning.

Neither should ever fire, and if one does it produces a silently wrong tensor rather than a loud failure.

Concretely, the whole visitor could be:

void DistributedCodegen::VisitStmt_(const ir::IfStmtPtr& op) {
  INTERNAL_CHECK(op != nullptr) << "Internal error: null IfStmt";

  // ConvertToSSA always appends an else carrying the phi's incoming values
  // (convert_to_ssa_pass.cpp:936-939), so a phi is always defined on both paths.
  INTERNAL_CHECK_SPAN(op->return_vars_.empty() || op->else_body_.has_value(), op->span_)
      << "Internal error: IfStmt with return_vars_ must carry an else_body "
         "holding the phi's incoming values";

  VisitExpr(op->condition_);
  const std::string condition = current_expr_value_;
  current_expr_value_ = "";

  // Merge each branch's yield into the phi name so post-if consumers see one
  // name instead of a branch-local SSA name (issue #2180).
  auto emit_yields = [&](const ir::StmtPtr& body) {
    const auto yld =
        ir::transform_utils::GetLastYieldStmt(ir::transform_utils::UnwrapAutoScope(body));
    if (!yld) return;
    for (size_t i = 0; i < op->return_vars_.size() && i < yld->value_.size(); ++i) {
      VisitExpr(yld->value_[i]);
      const std::string val = current_expr_value_;
      current_expr_value_ = "";
      const std::string phi = SanitizeName(op->return_vars_[i]->name_hint_);
      if (ir::AsTensorTypeLike(op->return_vars_[i]->GetType())) {
        emitter_.EmitLine("tensors[\"" + phi + "\"] = tensors[\"" + val + "\"]");
      } else {
        emitter_.EmitLine(phi + " = " + val);
      }
      declared_vars_.insert(phi);
    }
  };

  emitter_.EmitLine("if " + condition + ":");
  emitter_.IncreaseIndent();
  VisitStmt(op->then_body_);
  emit_yields(op->then_body_);
  emitter_.DecreaseIndent();

  if (op->else_body_.has_value()) {
    emitter_.EmitLine("else:");
    emitter_.IncreaseIndent();
    VisitStmt(*op->else_body_);
    emit_yields(*op->else_body_);
    emitter_.DecreaseIndent();
  }
}

This also merges the phi and non-phi paths (the return_vars_.empty() loop is a no-op), so there's no duplicated if/else emission, and it drops the transform_utils.h include down to just the two helpers actually used.

Side benefit: find_yield_source calls FlattenToStmts once per phi, and is called from both the init loop and the yield loop, so the current shape is O(P·N) per IfStmt (times nesting depth). .claude/rules/pass-complexity.md asks for O(N log N); the version above is a single linear pass.

2. The else_body_ check is not just defensive

Worth keeping the INTERNAL_CHECK_SPAN above, because the frontend can produce return_vars_ without an else. parse_if_statement (ast_parser.py:3050) creates return_vars_ from explicit pl.yield_() calls (if_builder.return_var(...), lines 3105–3113) but only calls if_builder.else_() when stmt.orelse exists (line 3096). So:

if r == 0:
    pl.yield_(boundary=zero)   # no else

yields IfStmt(return_vars_=[boundary], else_body_=nullopt). ConvertToSSA does not repair it either: since zero is a parameter that is never re-versioned inside the branch, then_ver == before and phis comes out empty (lines 843–860), which takes the pass-through path at lines 876–893 where new_else stays nullopt. The else synthesis at 936–939 is only reached when phis is non-empty.

Today that shape silently emits the (1,)/float32 placeholder; without the pre-declaration it would regress to the original KeyError. An explicit internal check turns it into a clear compiler-bug message. (I did not check whether a verifier already rejects this shape — if it does, the check is just cheap insurance.)

3. As<Var> misses IterArg in the AssignStmt branch

This one survives the simplification, since it's the core fix:

if (ir::AsTensorTypeLike(op->var_->GetType()) && ir::As<ir::Var>(op->value_)) {

As<T>() is an exact ObjectKind match, so IterArg doesn't hit it (see .claude/rules/ir-kind-traits.md). Loop-carried tensors from pl.range(..., init_values=[...]) are IterArgs, and they'd fall through to the old bare-name path — the same bug this PR fixes, via a different door. Suggest ir::AsVarLike(op->value_).

Pre-existing and out of scope, but related: DistributedCodegen overrides VisitExpr_(const ir::VarPtr&) rather than VisitVarLike_, while functor.h:111 dispatches IterArg to VisitExpr_(IterArgPtr) → the base IRVisitor::VisitVarLike_ (visitor.h:48-52), which never sets current_expr_value_. So even with AsVarLike, VisitExpr(iter_arg) returns an empty string. If loop-carried tensors are meant to work, VisitExpr_(VarPtr) should become VisitVarLike_; otherwise it's worth a follow-up issue.

4. The AsTensorTypeLike guard over-matches DistributedTensorType

AsTensorTypeLike matches both TensorType and DistributedTensorType (kind_traits.h:329). The distributed_tensor_alias early return above only fires when both sides are DistributedTensorType; a DistributedTensorType var with a plain-TensorType Var value would reach the new branch and emit tensors[x] = tensors[y], whereas DistributedTensors are resolved through window_buffer_ (lines 1005–1008), not the tensors dict. Exact ir::As<ir::TensorType>(op->var_->GetType()) here would be safer. (Distributed-tensor phis aren't handled by either the old or the new code — probably fine to leave, but worth a comment.)

5. Test

Good that it's a real before/after with assertions, in the right place, and that _lower(..., convert_to_ssa=False) keeps existing callers untouched. Two things:

• code.split("if", 1) splits on the substring if, which matches comments, identifiers, and — most relevant here — the if "<target>" not in tensors: line that EmitCallToWorker emits at line 1068 for a callee with no Out params, which is exactly what chip_run is in this test. It happens to pass today, but any preamble change can silently move the split point. A line-anchored re.search(r"^\s*if .*:$", line, re.M) would be robust.
• Since the reported symptom is "generated module fails at prepare()", a compile(code, "<generated>", "exec") assertion (or an ast.parse) covers the bug more directly than the name regexes.

If you keep the simplified version, a scalar-phi case and an if + pl.yield_() without else case would cover the two paths flagged above cheaply.

Overall: the fix is correct, CI is green, and the risk is confined to the distributed codegen path. My suggestion is to land the AssignStmt change plus the ~25-line IfStmt version, which removes the fallback initializers and the O(P·N) scan along with it.

Review

The diagnosis and the fix direction are right: in this codegen a tensor only ever exists as a tensors["..."] entry (every use site goes through it — e.g. EmitCallToWorker at distributed_codegen.cpp:1018), so a bare x = y Python assignment for a tensor alias is dead on arrival. The new branch in VisitStmt_(AssignStmtPtr) is the correct fix for #2180.

My main comment is that the ~140 new lines in VisitStmt_(IfStmtPtr) can collapse to ~25, and most of the remaining concerns disappear with it.

1. find_yield_source and the phi pre-declaration are redundant

The comment justifying find_yield_source says the then-branch's boundary = zero creates a bare Python name that isn't in the tensors dict. That was true before this PR — but the AssignStmt fix in this same PR removes that case. Walking every way a yield var can be defined:

definition of the yield var is tensors[<name>] populated?
Var→Var tensor alias ✅ the new AssignStmt branch in this PR
Call RHS, callee has Out/InOut params ✅ EmitCallToWorker, line 1130
Call RHS, no Out params ✅ EmitCallToWorker, lines 1068–1070
tensor.create ✅ EmitTensorCreate, line 1194
TupleGetItemExpr unpacking ✅ line 713
function parameter ✅ seeded by the runtime
hoisted alloc ✅ _alloc_intermediates already filled it
So by the time VisitStmt(op->then_body_) returns, tensors[<yield var name>] is guaranteed to exist and the yield can be emitted directly, with no RHS back-tracing.

The pre-declaration is redundant for the same reason: ConvertToSSA always synthesizes an else body when it creates phis (convert_to_ssa_pass.cpp:936-939, else_with_yield = make_shared<YieldStmt>(else_yields, ...)), so both branches always assign the phi. That also removes the two fallback initializers, which are the part I'd most like to see gone:

• torch.zeros((1,), dtype=torch.float32) hardcodes shape and dtype, while the phi's own TensorType is right there (compare lines 1051–1070, which derive shape/dtype from the return type).
• tensor_phi_init picks the function's first tensor parameter, which is unrelated to the phi in shape, dtype and meaning.

Neither should ever fire, and if one does it produces a silently wrong tensor rather than a loud failure.

Concretely, the whole visitor could be:

void DistributedCodegen::VisitStmt_(const ir::IfStmtPtr& op) {
  INTERNAL_CHECK(op != nullptr) << "Internal error: null IfStmt";

  // ConvertToSSA always appends an else carrying the phi's incoming values
  // (convert_to_ssa_pass.cpp:936-939), so a phi is always defined on both paths.
  INTERNAL_CHECK_SPAN(op->return_vars_.empty() || op->else_body_.has_value(), op->span_)
      << "Internal error: IfStmt with return_vars_ must carry an else_body "
         "holding the phi's incoming values";

  VisitExpr(op->condition_);
  const std::string condition = current_expr_value_;
  current_expr_value_ = "";

  // Merge each branch's yield into the phi name so post-if consumers see one
  // name instead of a branch-local SSA name (issue #2180).
  auto emit_yields = [&](const ir::StmtPtr& body) {
    const auto yld =
        ir::transform_utils::GetLastYieldStmt(ir::transform_utils::UnwrapAutoScope(body));
    if (!yld) return;
    for (size_t i = 0; i < op->return_vars_.size() && i < yld->value_.size(); ++i) {
      VisitExpr(yld->value_[i]);
      const std::string val = current_expr_value_;
      current_expr_value_ = "";
      const std::string phi = SanitizeName(op->return_vars_[i]->name_hint_);
      if (ir::AsTensorTypeLike(op->return_vars_[i]->GetType())) {
        emitter_.EmitLine("tensors[\"" + phi + "\"] = tensors[\"" + val + "\"]");
      } else {
        emitter_.EmitLine(phi + " = " + val);
      }
      declared_vars_.insert(phi);
    }
  };

  emitter_.EmitLine("if " + condition + ":");
  emitter_.IncreaseIndent();
  VisitStmt(op->then_body_);
  emit_yields(op->then_body_);
  emitter_.DecreaseIndent();

  if (op->else_body_.has_value()) {
    emitter_.EmitLine("else:");
    emitter_.IncreaseIndent();
    VisitStmt(*op->else_body_);
    emit_yields(*op->else_body_);
    emitter_.DecreaseIndent();
  }
}

This also merges the phi and non-phi paths (the return_vars_.empty() loop is a no-op), so there's no duplicated if/else emission, and it drops the transform_utils.h include down to just the two helpers actually used.

Side benefit: find_yield_source calls FlattenToStmts once per phi, and is called from both the init loop and the yield loop, so the current shape is O(P·N) per IfStmt (times nesting depth). .claude/rules/pass-complexity.md asks for O(N log N); the version above is a single linear pass.

2. The else_body_ check is not just defensive

Worth keeping the INTERNAL_CHECK_SPAN above, because the frontend can produce return_vars_ without an else. parse_if_statement (ast_parser.py:3050) creates return_vars_ from explicit pl.yield_() calls (if_builder.return_var(...), lines 3105–3113) but only calls if_builder.else_() when stmt.orelse exists (line 3096). So:

if r == 0:
    pl.yield_(boundary=zero)   # no else

yields IfStmt(return_vars_=[boundary], else_body_=nullopt). ConvertToSSA does not repair it either: since zero is a parameter that is never re-versioned inside the branch, then_ver == before and phis comes out empty (lines 843–860), which takes the pass-through path at lines 876–893 where new_else stays nullopt. The else synthesis at 936–939 is only reached when phis is non-empty.

Today that shape silently emits the (1,)/float32 placeholder; without the pre-declaration it would regress to the original KeyError. An explicit internal check turns it into a clear compiler-bug message. (I did not check whether a verifier already rejects this shape — if it does, the check is just cheap insurance.)

3. As<Var> misses IterArg in the AssignStmt branch

This one survives the simplification, since it's the core fix:

if (ir::AsTensorTypeLike(op->var_->GetType()) && ir::As<ir::Var>(op->value_)) {

As<T>() is an exact ObjectKind match, so IterArg doesn't hit it (see .claude/rules/ir-kind-traits.md). Loop-carried tensors from pl.range(..., init_values=[...]) are IterArgs, and they'd fall through to the old bare-name path — the same bug this PR fixes, via a different door. Suggest ir::AsVarLike(op->value_).

Pre-existing and out of scope, but related: DistributedCodegen overrides VisitExpr_(const ir::VarPtr&) rather than VisitVarLike_, while functor.h:111 dispatches IterArg to VisitExpr_(IterArgPtr) → the base IRVisitor::VisitVarLike_ (visitor.h:48-52), which never sets current_expr_value_. So even with AsVarLike, VisitExpr(iter_arg) returns an empty string. If loop-carried tensors are meant to work, VisitExpr_(VarPtr) should become VisitVarLike_; otherwise it's worth a follow-up issue.

4. The AsTensorTypeLike guard over-matches DistributedTensorType

AsTensorTypeLike matches both TensorType and DistributedTensorType (kind_traits.h:329). The distributed_tensor_alias early return above only fires when both sides are DistributedTensorType; a DistributedTensorType var with a plain-TensorType Var value would reach the new branch and emit tensors[x] = tensors[y], whereas DistributedTensors are resolved through window_buffer_ (lines 1005–1008), not the tensors dict. Exact ir::As<ir::TensorType>(op->var_->GetType()) here would be safer. (Distributed-tensor phis aren't handled by either the old or the new code — probably fine to leave, but worth a comment.)

5. Test

Good that it's a real before/after with assertions, in the right place, and that _lower(..., convert_to_ssa=False) keeps existing callers untouched. Two things:

• code.split("if", 1) splits on the substring if, which matches comments, identifiers, and — most relevant here — the if "<target>" not in tensors: line that EmitCallToWorker emits at line 1068 for a callee with no Out params, which is exactly what chip_run is in this test. It happens to pass today, but any preamble change can silently move the split point. A line-anchored re.search(r"^\s*if .*:$", line, re.M) would be robust.
• Since the reported symptom is "generated module fails at prepare()", a compile(code, "<generated>", "exec") assertion (or an ast.parse) covers the bug more directly than the name regexes.

If you keep the simplified version, a scalar-phi case and an if + pl.yield_() without else case would cover the two paths flagged above cheaply.

Overall: the fix is correct, CI is green, and the risk is confined to the distributed codegen path. My suggestion is to land the AssignStmt change plus the ~25-line IfStmt version, which removes the fallback initializers and the O(P·N) scan along with it.

I can't authenticate gh to post directly, but here's the reply formatted to copy-paste into the PR as a response to YunjiQin's review:

@YunjiQin thanks for the detailed review. Addressed all points in the latest push:

1. Simplified VisitStmt_(IfStmtPtr) from ~140 lines down to ~40
Dropped find_yield_source, the pre-declaration loop, and both fallback initializers (torch.zeros(…) and tensor_phi_init). With the AssignStmt fix in place, tensors[…] is always populated by the time VisitStmt returns, so yields can be emitted directly — no RHS back-tracing needed. The return_vars_ / else_body_ paths are unified under a single emit_yields lambda (no-op when return_vars_ is empty). Added the INTERNAL_CHECK_SPAN you suggested — consistent with the same guard already in PTO codegen, the SSA verifier, and the simplify pass.

2. AsVarLike instead of As<Var> in the AssignStmt branch
Loop-carried IterArg tensors now hit the tensors[…] path instead of the bare-name fallthrough.

3. Exact As<TensorType> instead of AsTensorTypeLike in the AssignStmt branch
DistributedTensorType won't reach the new branch — distributed tensors go through window_buffer_, not the tensors dict.

4. Test robustness
Replaced code.split("if", 1) with line-anchored ^\s*if\s+.*:\s*$ regex and indent-aware else: matching. Added compile(code, "<host_orch>", "exec") as a syntactic sanity check. Renamed to test_if_cross_branch_phi_yields_tensors since there's no pre-declaration anymore.

All 8015 unit tests, 637 codegen tests, and 125 SSA tests pass. No existing test exercises the pl.yield_()-without-else path — ConvertToSSA always synthesizes an else_body_ when it creates phis, so the INTERNAL_CHECK_SPAN is purely defensive insurance.